#!/usr/bin/python

import numpy,os,os.path,sys
from collections import Counter
from pylab import *
import tables,scipy
from scipy import fftpack
from scipy.cluster import vq
from scipy.ndimage import morphology,filters,interpolation
from ocrolib import docproc,improc
import pyflann
import random
from tables import openFile,Filters,Int32Atom,Float32Atom,Int64Atom
from multiprocessing import Pool
import fcntl
from sutils import *
from ocrolib.pycomp import PyComponent
from pdb import pm

def sigmoid(x):
    return 1/(1+exp(-clip(x,-20,20)))

class LocalLinearModel(PyComponent):
    def __init__(self,**kw):
        self.fname = "kmeans-500s.h5"
        self.rescore = None
        self.dilation = 0
        self.floor = 0.1
        self.costfactor = 1.0
    def get_ll(self,j):
        i = self.ll_which.get(j,None)
        if i is None: return None
        pdb = self.pdb
        classes = array(pdb.root.ll_classes[i],'i')
        n = find(classes>0)[-1]+1
        classes = array(classes[:n])
        ms = array(pdb.root.ll_ms[i,:n],'f')
        bs = array(pdb.root.ll_bs[i,:n],'f')
        ds = pdb.root.ll_ds[i]
        return (classes,ms,bs,ds)
    def load(self,pfile):
        print "loading"
        with openFile(pfile) as pdb:
            nprotos = len(pdb.root.patches)
            self.nprotos = nprotos
            self.protos = array(pdb.root.patches[:nprotos,:,:])
            self.pclasses = array(pdb.root.classes[:nprotos])
            self.pcounts = array(pdb.root.counts[:nprotos])
            self.psigmas = array(pdb.root.sigmas[:nprotos,:,:])
            self.classes = array(pdb.root.classes[:nprotos])
            if "relmeans" in dir(pdb.root):
                self.relmeans = array(pdb.root.relmeans[:nprotos,:])
                self.relsigmas = array(pdb.root.relsigmas[:nprotos,:])
                print "loaded relmeans"
            else:
                print "warning: no rel information"
            if "ll_which" in dir(pdb.root):
                self.ll_which = {}
                which = array(pdb.root.ll_which[:len(pdb.root.ll_which)])
                for i,w in enumerate(which):
                    self.ll_which[w] = i
                print "loaded ll information"
            else:
                ll_which = None
                print "warning: no ll information"
            if "nn1" in dir(pdb.root):
                self.nn1 = array(pdb.root.nn1[:nprotos,:,:])
            else:
                self.nn1 = None
            if "nnstats" in dir(pdb.root):
                self.nnstats = array(pdb.root.nnstats[:nprotos,:,:])
            else:
                self.nnstats = None
        self.pdb = openFile(pfile)
        self.nnstats = None
        self.pindex = NNIndex()
        self.pindex.build_index(make2d(self.protos))
        print "done"
    def set(self,**kw):
        pass
    def linefilter(self,rel,results,dists):
        for i in range(len(results)):
            p = results[i]
            rmeans = self.relmeans[p]
            rsigmas = self.relsigmas[p]
            deltas = (rel-rmeans)**2/(2*maximum(0.01,rsigmas)**2)
            # print rel,rmeans,rsigmas,deltas
            if args.linefilter<0:
                print amax(deltas)
            elif (deltas>args.linefilter).any():
                dists[i] = 999.0
        order = argsort(dists)
        return results[order],dists[order]
    def clear(self):
        raise Exception("unimplemented")
    def cadd(self,image,c,geometry=None):
        raise Exception("unimplemented")
    def updateModel(self,*args,**kw):
        raise Exception("unimplemented")
    def coutputs(self,image,rel=None):
        if self.protos is None: self.load(self.fname)
        # find candidate nearest neighbors
        [presult],[pdists] = self.pindex.nn_index(array([image.ravel()]),maximum(2,self.rescore))
        # perform the rescoring
        if self.rescore>0:
            raise Exception("unimplemented")
            # FIXME
            order = argsort(pdists)
            presult = presult[order]
            pdists = pdists[order]
        # apply the line filter
        if rel is not None:
            presult,pdists = self.linefilter(rel,presult,pdists)
            order = argsort(pdists)
            presult = presult[order]
            pdists = pdists[order]
        n1 = presult[0]
        d1 = pdists[0]
        # compute rejection, using the distribution of distances
        cost = 0
        if self.nnstats is not None:
            mu = self.nnstats[n1,1,0]
            sigma = self.nnstats[n1,2,0]
            mu_r = self.nnstats[n1,1,2]
            sigma_r = self.nnstats[n1,2,2]
            cost = 0.5*((max(mu,d1)-mu)/sigma)**2
            # cost_r = 0.5*((min(mu_r,d1)-mu_r)/sigma_r)**2
            # print d1,mu,sigma,cost,cost_r
            cost = clip(cost,0,9)
        self.last = (n1,d1)
        if self.ll_which is not None:
            lin = cm.get_ll(n1)
            if lin is None:
                result = [(udecode(self.classes[n1]),0.999)]
                return result
            classes,ms,bs,ds = lin
            out = sigmoid(dot(ms,image.ravel())+bs)
            result = zip([udecode(c) for c in classes],out)
            result = sorted(result,key=lambda x:-x[1])
            return result
        if self.nn1 is None:
            result = [(udecode(self.classes[n1]),0.999)]
            return result
        result = self.nn1[n1]
        result = [(udecode(c),n) for c,n in result if c>0 and c!=126 and c!=ord('_')]
        total = sum([n for c,n in result])
        result = [(c,exp(-cost)*n*1.0/total) for c,n in result]
        if cost>1e-6: result += [("~",1.0-exp(-cost))]
        result = sorted(result,key=lambda x:-x[1])
        return result
    def cclassify(self,v,geometry=None):
        pass
    def coutputs_batch(self,images,geometries=None):
        result = []
        for i in range(len(images)):
            try:
                output = self.coutputs(images[i],geometries[i]) 
            except:
                print "recognition failed"
                output = []
            result.append(output)
        return result
    def save_component(self,path):
        pass


db = tables.openFile("data/dtraining.h5")
classes = array(db.root.classes[:10000])
image = db.root.patches[find(array(classes!=126))[0]]
cm = LocalLinearModel()
cm.load("nnl.h5")

for i in range(1000,1100):
    cls = udecode(db.root.classes[i])
    image = db.root.patches[i]
    print i,cls,cm.coutputs(image)
